"""
NVIDIA显卡动态状态信息采集脚本
采集显卡实时负载、温度、功耗、显存使用等动态运行状态
"""

import subprocess
import json
from typing import List, Optional


class NVIDIAGPUDynamicCollector:
    def get_gpu_dynamic_info(self, gpu_index: int) -> Optional[dict]:
        """获取单个GPU的动态运行状态"""
        try:
            # 查询GPU动态信息
            cmd = [
                'nvidia-smi', '-i', str(gpu_index),
                '--query-gpu=name,memory.used,memory.free,utilization.gpu,utilization.memory,temperature.gpu,power.draw,power.limit,clocks.gr,clocks.mem',
                '--format=csv,noheader,nounits'
            ]

            result = subprocess.run(cmd, capture_output=True, text=True, check=True)
            data = result.stdout.strip().split(', ')

            if len(data) < 10:
                return None

            # 安全转换数值
            def safe_int(value):
                try:
                    return int(value) if value != '[Not Supported]' else 0
                except:
                    return 0

            def safe_float(value):
                try:
                    return float(value) if value != '[Not Supported]' else 0.0
                except:
                    return 0.0

            memory_used = safe_int(data[1])
            memory_free = safe_int(data[2])
            memory_total = memory_used + memory_free
            gpu_util = safe_int(data[3])
            memory_util = safe_int(data[4])
            power_draw = safe_float(data[6])
            power_limit = safe_float(data[7])

            gpu_info = {
                'index': gpu_index,  # GPU索引
                'name': data[0],  # GPU名称
                'memory_usage': {  # 显存使用情况
                    'total_mb': memory_total,  # 总显存MB
                    'used_mb': memory_used,  # 已用显存MB
                    'free_mb': memory_free,  # 空闲显存MB
                    'usage_percent': round((memory_used / memory_total * 100) if memory_total > 0 else 0, 1)  # 使用率%
                },
                'workload': {  # 工作负载
                    'gpu_utilization_percent': gpu_util,  # GPU核心利用率%
                    'memory_utilization_percent': memory_util,  # 显存带宽利用率%
                    'current_gpu_clock_mhz': safe_int(data[8]),  # 当前GPU频率MHz
                    'current_memory_clock_mhz': safe_int(data[9])  # 当前显存频率MHz
                },
                'thermal_power': {  # 温度功耗
                    'temperature_celsius': safe_int(data[5]),  # 温度℃
                    'power_draw_watts': power_draw,  # 当前功耗W
                    'power_limit_watts': power_limit,  # 功耗限制W
                    'power_usage_percent': round((power_draw / power_limit * 100) if power_limit > 0 else 0, 1)
                    # 功耗使用率%
                }
            }

            return gpu_info

        except subprocess.CalledProcessError:
            return None

    def calculate_dynamic_summary(self, gpus: List[dict]) -> dict:
        """计算动态信息汇总统计"""
        if not gpus:
            return {
                'gpu_count': 0,
                'memory_summary': {
                    'total_mb': 0,
                    'used_mb': 0,
                    'free_mb': 0,
                    'usage_percent': 0.0
                },
                'workload_summary': {
                    'avg_gpu_utilization_percent': 0.0,
                    'avg_memory_utilization_percent': 0.0,
                    'max_gpu_utilization_percent': 0,
                    'min_gpu_utilization_percent': 0
                },
                'thermal_power_summary': {
                    'avg_temperature_celsius': 0.0,
                    'max_temperature_celsius': 0,
                    'total_power_draw_watts': 0.0,
                    'total_power_limit_watts': 0.0,
                    'power_usage_percent': 0.0
                }
            }

        # 显存统计
        total_memory_mb = sum(gpu['memory_usage']['total_mb'] for gpu in gpus)
        total_memory_used_mb = sum(gpu['memory_usage']['used_mb'] for gpu in gpus)
        total_memory_free_mb = sum(gpu['memory_usage']['free_mb'] for gpu in gpus)

        # 负载统计
        avg_gpu_utilization = sum(gpu['workload']['gpu_utilization_percent'] for gpu in gpus) / len(gpus)
        avg_memory_utilization = sum(gpu['workload']['memory_utilization_percent'] for gpu in gpus) / len(gpus)
        max_gpu_utilization = max(gpu['workload']['gpu_utilization_percent'] for gpu in gpus)
        min_gpu_utilization = min(gpu['workload']['gpu_utilization_percent'] for gpu in gpus)

        # 温度功耗统计
        avg_temperature = sum(gpu['thermal_power']['temperature_celsius'] for gpu in gpus) / len(gpus)
        max_temperature = max(gpu['thermal_power']['temperature_celsius'] for gpu in gpus)
        total_power_draw = sum(gpu['thermal_power']['power_draw_watts'] for gpu in gpus)
        total_power_limit = sum(gpu['thermal_power']['power_limit_watts'] for gpu in gpus)

        return {
            'gpu_count': len(gpus),  # GPU数量
            'memory_summary': {  # 显存汇总
                'total_mb': total_memory_mb,  # 总显存MB
                'used_mb': total_memory_used_mb,  # 已用显存MB
                'free_mb': total_memory_free_mb,  # 空闲显存MB
                'usage_percent': round((total_memory_used_mb / total_memory_mb * 100) if total_memory_mb > 0 else 0, 1)
                # 整体显存使用率%
            },
            'workload_summary': {  # 负载汇总
                'avg_gpu_utilization_percent': round(avg_gpu_utilization, 1),  # 平均GPU利用率%
                'avg_memory_utilization_percent': round(avg_memory_utilization, 1),  # 平均显存带宽利用率%
                'max_gpu_utilization_percent': max_gpu_utilization,  # 最高GPU利用率%
                'min_gpu_utilization_percent': min_gpu_utilization  # 最低GPU利用率%
            },
            'thermal_power_summary': {  # 温度功耗汇总
                'avg_temperature_celsius': round(avg_temperature, 1),  # 平均温度℃
                'max_temperature_celsius': max_temperature,  # 最高温度℃
                'total_power_draw_watts': round(total_power_draw, 1),  # 总功耗W
                'total_power_limit_watts': round(total_power_limit, 1),  # 总功耗限制W
                'power_usage_percent': round(
                    (total_power_draw / total_power_limit * 100) if total_power_limit > 0 else 0, 1)  # 整体功耗使用率%
            }
        }

    def collect_dynamic_info(self) -> dict:
        """收集所有GPU动态信息"""
        try:
            # 检查nvidia-smi是否可用
            subprocess.run(['nvidia-smi', '-L'], capture_output=True, check=True)

            # 获取GPU列表
            result = subprocess.run(['nvidia-smi', '-L'], capture_output=True, text=True, check=True)
            gpu_lines = result.stdout.strip().split('\n')

            if not gpu_lines or gpu_lines == ['']:
                return {
                    'status': 'no_gpu',  # 状态：无GPU
                    'message': '未检测到NVIDIA显卡',
                    'summary': self.calculate_dynamic_summary([]),
                    'gpus': []
                }

            # 获取所有GPU动态信息
            gpus = []
            for i in range(len(gpu_lines)):
                gpu_info = self.get_gpu_dynamic_info(i)
                if gpu_info:
                    gpus.append(gpu_info)

            # 计算汇总统计
            summary = self.calculate_dynamic_summary(gpus)

            return {
                'status': 'success',  # 状态：成功
                'message': f'成功获取 {len(gpus)} 张显卡动态信息',
                'summary': summary,  # 汇总信息
                'gpus': gpus  # 详细GPU列表
            }

        except (subprocess.CalledProcessError, FileNotFoundError):
            return {
                'status': 'no_driver',  # 状态：无驱动
                'message': 'NVIDIA驱动未安装或nvidia-smi不可用',
                'summary': self.calculate_dynamic_summary([]),
                'gpus': []
            }

def main():
    collector = NVIDIAGPUDynamicCollector()
    result = collector.collect_dynamic_info()
    print(json.dumps(result, indent=2, ensure_ascii=False))

if __name__ == "__main__":
    main()

